# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import collections
import copy
import functools
import logging
import os

import numpy
import pandas
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages

from coin.base.book.types import TradeSide
from coin.base.datetime_util import to_timestamp_int

import coin.strategy.mm.tool.archive_base as abase

from coin2.support.product_info.product_info_holder import ProductInfoHolder


def get_marker(compare_idx):
  return f"C{compare_idx}-"


def combine_names(names):
  return "_".join([str(name) for name in names if name is not None])


class FeedArchiveChecker(object):
  def __init__(self,
               products,
               start_ts=None,
               end_ts=None,
               product_to_sub_request=None,
               use_exchange_timestamp=None):
    self._products = products
    self._tick_sizes = (0.5, 0.01)
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)
    self._fields_out = []
    self._fields_view = []
    self._time_set = set()
    self._agg_book = {}
    self._bbos = {product: [] for product in products}
    self._trades = {product: [] for product in products}
    self._instrs = {product: [] for product in products}
    self.nbook = 0
    self.ntrade = 0
    self.product_to_sub_request = product_to_sub_request
    self.use_exchange_timestamp = use_exchange_timestamp

  def report(self):
    for product, bbo in self._bbos.items():
      print(ProductInfoHolder.relative_norm(product), " # feed: ", len(bbo))

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))
      # if hasattr(book_builder, "subscribe_instrument"):
      #  book_builder.subscribe_instrument(
      #    product, functools.partial(self.on_instrument, product))

  def on_instrument(self, product, instrument):
    if hasattr(instrument, 'open_interest'):
      open_interest = instrument.open_interest
      funding_rate = instrument.funding_rate
      mark_price = instrument.mark_price
      indicative_funding_rate = instrument.indicative_funding_rate
    else:
      return

    self._instrs[product].append({
        'timestamp': instrument.timestamp,
        'open_interest': open_interest,
        'funding_rate': funding_rate,
        'indicative_funding_rate': indicative_funding_rate,
        'mark_price': mark_price,
    })

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return

    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
    except (AttributeError, IndexError):
      return

    if flags.FLAGS.print_all_feed:
      askarr = book.get_ask_array(2)
      bidarr = book.get_bid_array(2)
      print("%d %.4f %.4f %.4f %.4f %.4f %.4f %.4f %.4f" % (book.timestamp,
                                                              askarr[0][0],
                                                              askarr[0][1],
                                                              askarr[1][0],
                                                              askarr[1][1],
                                                              bidarr[0][0],
                                                              bidarr[0][1],
                                                              bidarr[1][0],
                                                              bidarr[1][1]))

    self._bbos[product].append({
        'timestamp':
            book.exchange_timestamp if self.use_exchange_timestamp else book.timestamp,
        'bid0p':
            bid0_price,
        'ask0p':
            ask0_price
    })
    self.nbook += 1

  def on_trade(self, product, trade):
    if trade.timestamp < self._start_ts:
      return

    if self._end_ts and self._end_ts < trade.timestamp:
      return

    if flags.FLAGS.print_all_feed:
      print("%.0f %.4f %.4f" % (trade.timestamp, trade.price, trade.qty))

    self._trades[product].append({
        'timestamp':
            trade.exchange_timestamp if self.use_exchange_timestamp else trade.timestamp,
        'tradep':
            trade.price,
        'side':
            trade.side
    })
    self.ntrade += 1

  def plot_all(self,
               out_dir,
               postfix,
               trading_date,
               machine,
               compare_idx,
               excel_file=None,
               rasterized=None,
               multiday=False):
    if compare_idx is not None:
      assert len(self._bbos.keys()) == 1, "multiple symbol is not allowed in compare mode"
    legends = []
    for i, (product, bbo) in enumerate(self._bbos.items()):
      if excel_file is not None and i == 0:
        self.plot_excel_file(product, excel_file)
      legends.append(
          self.plot_product(trading_date,
                            product,
                            machine + postfix,
                            out_dir,
                            compare_idx,
                            rasterized=rasterized,
                            multiday=multiday))
    return legends

  def plot_excel_file(self, product, excel_file):
    exceldata1 = pandas.read_excel(excel_file, sheet_name='LimitOrder-1', header=1)
    exceldata2 = pandas.read_excel(excel_file, sheet_name='LimitOrder-2', header=1)

    df = exceldata1.append(exceldata2, ignore_index=True)

    cols = [
        column for column in df.columns
        if column.find('Cont') >= 0 or column.find('PL') >= 0 or column.find('Fee') >= 0
    ]
    df2 = df[cols]

    print(df2.groupby('Contracts').sum())

    df = df.loc[df['Contracts'] == product.native_symbol].reset_index(drop=True)
    df = df[::-1]

    df['price'] = df['Price($)']
    df['fill_qty'] = df['Filled(Cont)']
    df['qty'] = df['Amount(Cont)']
    df['side'] = ((df['Type'].str.find("Long") >= 0) - 0.5) * 2 * (
        (df['Type'].str.find("Open") >= 0) - 0.5) * 2
    df['time'] = pandas.DatetimeIndex(df['Date']).astype(int) - 8 * 3600 * 1e9
    bdf = df.loc[df['side'] == 1]
    sdf = df.loc[df['side'] == -1]
    plt.plot(pandas.DatetimeIndex(bdf['time']), bdf['price'], 'k:')
    plt.plot(pandas.DatetimeIndex(sdf['time']), sdf['price'], 'k:')
    tdf = df.loc[df['Price type'] == 'IOC']
    btdf = tdf.loc[tdf['side'] == 1]
    stdf = tdf.loc[tdf['side'] == -1]
    plt.plot(pandas.DatetimeIndex(btdf['time']), btdf['price'], 'g.', markersize=10)
    plt.plot(pandas.DatetimeIndex(stdf['time']), stdf['price'], 'r.', markersize=10)

  def plot_product(self,
                   date_str,
                   product,
                   machine,
                   out_dir,
                   compare_idx,
                   rasterized=None,
                   multiday=False):
    assert (out_dir is None or len(out_dir) > 0)
    assert product in self._products, (product, self._prdoucts)
    assert len(self._bbos[product]) > 0, product
    df = pandas.DataFrame(self._bbos[product])
    timeindex = pandas.DatetimeIndex(df['timestamp'])

    tdf = pandas.DataFrame(self._trades[product])
    line = None
    if compare_idx is not None:
      line = plt.plot(timeindex,
                      df['bid0p'],
                      get_marker(compare_idx),
                      timeindex,
                      df['ask0p'],
                      get_marker(compare_idx),
                      lw=0.2,
                      markersize=0.2,
                      drawstyle='steps-post',
                      rasterized=rasterized)
    else:
      if len(tdf) > 0:
        tindex = pandas.DatetimeIndex(tdf['timestamp'])
        idx_buy = tdf['side'] == TradeSide.TRADE_BUY_SIDE
        idx_sell = tdf['side'] == TradeSide.TRADE_SELL_SIDE
        idx_else = (~idx_buy) & (~idx_sell)

        plt.plot(tindex[idx_buy],
                 tdf['tradep'][idx_buy],
                 'gx',
                 tindex[idx_sell],
                 tdf['tradep'][idx_sell],
                 'rx',
                 tindex[idx_else],
                 tdf['tradep'][idx_else],
                 'mx',
                 timeindex,
                 df['bid0p'],
                 'k-',
                 timeindex,
                 df['ask0p'],
                 'k-',
                 lw=0.2,
                 markersize=2,
                 drawstyle='steps-post',
                 rasterized=rasterized)
        plt.xticks(rotation=90)
      else:
        plt.plot(timeindex, df['bid0p'], 'g-', timeindex, df['ask0p'], 'r-', lw=0.2, markersize=0.2)

    plt.minorticks_on()
    plt.grid(lw=0.5, which='major')
    plt.grid(lw=0.2, which='minor')

    if self._start_ts is not None and self._end_ts is not None:
      send = pandas.DatetimeIndex([self._start_ts, self._end_ts])
      plt.xlim(send[0], send[1])

    plt.minorticks_on()
    plt.grid(lw=0.5, which='major')
    plt.grid(lw=0.2, which='minor')

    oidf = pandas.DataFrame(self._instrs[product])

    if compare_idx is None and len(oidf) > 0:
      tindex = pandas.DatetimeIndex(oidf['timestamp'])
      nnmp = ~numpy.isnan(oidf['mark_price'])
      plt.plot(tindex[nnmp], oidf['mark_price'][nnmp], 'y:', lw=0.5, rasterized=rasterized)

      nnfr = ~numpy.isnan(oidf['indicative_funding_rate'])
      plt.twinx()
      plt.plot(tindex[nnfr],
               oidf['indicative_funding_rate'][nnfr],
               'g:',
               lw=0.5,
               rasterized=rasterized)
      plt.axis('off')

      nnoi = ~numpy.isnan(oidf['open_interest'])
      plt.twinx()
      plt.plot(tindex[nnoi], oidf['open_interest'][nnoi], 'k:', lw=0.5, rasterized=rasterized)

    os.makedirs(flags.FLAGS.out_dir, exist_ok=True)
    title = combine_names([
        date_str,
        machine,
        product.exchange,
        ProductInfoHolder.relative_norm(product),
        self.product_to_sub_request[product],
        flags.FLAGS.postfix,
        "ET" if self.use_exchange_timestamp else "FT"
    ])

    stat = "#B: %s #T: %s" % (self.nbook, self.ntrade)
    if compare_idx is not None or multiday:
      if multiday:
        plt.title("%s %s" % (title, stat))
      return (line, "%s %s" % (title, stat), (self.nbook, self.ntrade))
    else:
      filepath = os.path.join(out_dir, "%s.png" % title)
      plt.title("%s %s" % (title, stat))
      plt.savefig(filepath)
      plt.close()
      return None


def get_feedcache_machines():
  return [
      "feed-01.ap-northeast-1.aws",
      "feed-01.ap-northeast-2.aws",
      "feed-01.cn-hongkong.aliyun",
      "feed-01.eu-west-1.aws",
      "feed-01.us-east-1.aws",
      "feed-05.ap-northeast-1.aws",
  ]


def launch(baseflags,
           fsflags,
           symbols,
           out_dir,
           multiday,
           use_exchange_timestamp=None,
           *,
           compare_idx=None,
           rasterized=None,
           prerun=False):
  plt.rcParams['figure.figsize'] = 12, 8
  plt.rcParams['font.size'] = 9
  plt.rcParams['legend.fontsize'] = 9
  plt.rcParams['xtick.labelsize'] = 9
  plt.rcParams['ytick.labelsize'] = 9

  flags.FLAGS.use_exchange_timestamp = use_exchange_timestamp

  products = abase.symbols_to_products(symbols, baseflags)
  start_time, end_time = abase.get_time_range(baseflags)

  feedsub_request = abase.products_to_feedsub_request_with_api_override(
      products, fsflags.api_override, fsflags.feed_machine)

  product_to_sub_request = {}
  for sub_request, products in feedsub_request.gen_iter():
    for product in products:
      product_to_sub_request[product] = sub_request

  strategy = FeedArchiveChecker(products=products,
                                start_ts=start_time,
                                end_ts=end_time,
                                product_to_sub_request=product_to_sub_request,
                                use_exchange_timestamp=use_exchange_timestamp)

  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=products,
                              on_book_reset=strategy.on_book_reset)

  print(baseflags)
  strategy.report()

  kwargs = dict(out_dir=out_dir,
                postfix="wid_%s" % "_".join(fsflags.worker_ids.split(",")),
                trading_date=baseflags.trading_date,
                machine=fsflags.feed_machine,
                compare_idx=compare_idx,
                excel_file=flags.FLAGS.excel_file,
                rasterized=rasterized,
                multiday=multiday)
  return_args = (strategy.plot_all, kwargs)

  if prerun:
    return return_args
  else:
    return_args[0](**return_args[1])


def main(argv):
  trading_date_given = flags.FLAGS.trading_date

  fsflags = abase.get_feed_subscription_flags()

  fs_overrides = [None]

  if flags.FLAGS.full_symbols is not None:
    fs_overrides = list(abase.get_flags_with_full_symbols(flags.FLAGS.full_symbols))

  if flags.FLAGS.subsymbols is not None:
    fs_overrides = list(abase.get_flags_with_full_symbols(flags.FLAGS.subsymbols))

  out_dir = flags.FLAGS.out_dir
  assert out_dir
  trading_dates = abase.get_trading_dates(flags.FLAGS.trading_date)

  params = []

  if flags.FLAGS.compare_subsymbols is not None:
    compare = True
    fs_overrides = list(abase.get_flags_with_full_symbols(flags.FLAGS.compare_subsymbols))
  else:
    compare = False

  time_range = flags.FLAGS.time_range
  compare_groups = collections.OrderedDict()

  def generate_overrides():
    compare = False
    if flags.FLAGS.compare_use_exchange_timestamp is not None:
      use_exchange_timestamp_overrides = [
          bool(int(val)) for val in flags.FLAGS.compare_use_exchange_timestamp.split(",")
      ]
      compare = True
    else:
      use_exchange_timestamp_overrides = [flags.FLAGS.use_exchange_timestamp]

    if flags.FLAGS.compare_worker_ids is not None:
      worker_ids_overrides = flags.FLAGS.compare_worker_ids.split("-")
      compare = True
    else:
      worker_ids_overrides = [flags.FLAGS.worker_ids]

    if flags.FLAGS.compare_feed_machine_api_overrides is not None:
      feed_machine_api_overrides = flags.FLAGS.compare_feed_machine_api_overrides.split(",")
      compare = True

      for feed_machine_api_override in feed_machine_api_overrides:
        feed_machine, api_override = feed_machine_api_override.split("^")
        for worker_ids_override in worker_ids_overrides:
          for use_exchange_timestamp in use_exchange_timestamp_overrides:
            yield compare, {
                'worker_ids_override': worker_ids_override,
                'api_override': api_override,
                'feed_machine': feed_machine,
                'use_exchange_timestamp': use_exchange_timestamp
            }
    else:
      if flags.FLAGS.compare_api_overrides is not None:
        api_overrides = flags.FLAGS.compare_api_overrides.split(",")
        compare = True
      else:
        api_overrides = [flags.FLAGS.api_override]

      if flags.FLAGS.compare_feed_machines is not None:
        feed_machines = flags.FLAGS.compare_feed_machines.split(",")
        compare = True
      else:
        feed_machines = [fsflags.feed_machine]

      for worker_ids_override in worker_ids_overrides:
        for api_override in api_overrides:
          for feed_machine in feed_machines:
            for use_exchange_timestamp in use_exchange_timestamp_overrides:
              yield compare, {
                  'worker_ids_override': worker_ids_override,
                  'api_override': api_override,
                  'feed_machine': feed_machine,
                  'use_exchange_timestamp': use_exchange_timestamp
              }

  def find_three_option_files(options_dir):
    rtn_option_files = []
    for subdir, _, files in os.walk(options_dir):
      for file in files:
        if os.stat(os.path.join(subdir, file)).st_size > 10240:
          rtn_option_files.append(file.replace('.W0.lz4', ''))
        if len(rtn_option_files) > 3:
          return rtn_option_files
    return rtn_option_files

  def convert_options_symbols(symbols, fsflags):
    rtn_symbols = []
    for k, v in fsflags.api_override.items():
      options_dir = "/remote/iosg/coin-2/feed_cache2/%s/%s/%s.%s/" % (
          flags.FLAGS.feed_machine, flags.FLAGS.trading_date, k.replace(":", "."), v)
      for symbol in symbols:
        if "RANDOM" in symbol.upper():
          option_files = find_three_option_files(options_dir)
          if len(option_files) > 0:
            rtn_symbols.extend(option_files)
        else:
          rtn_symbols.append(symbol)
    return rtn_symbols

  for trading_date in trading_dates:
    flags.FLAGS.trading_date = trading_date.strftime("%Y%m%d")
    for trange in abase.split_time_range(
        time_range,
        flags.FLAGS.split_hours) if flags.FLAGS.split_hours is not None else [time_range]:
      cgroups = []
      for fs_override in fs_overrides:
        abase.set_fs_dict(fs_override)
        symbols = flags.FLAGS.symbols.split(",")

        baseflags = abase.get_base_flags()
        baseflags.time_range = trange
        for will_compare, fs_dict in generate_overrides():
          compare = compare or will_compare
          worker_ids_override = fs_dict['worker_ids_override']
          api_override = fs_dict['api_override']
          feed_machine = fs_dict['feed_machine']

          fsflags.worker_ids = worker_ids_override
          fsflags.feed_machine = feed_machine
          fsflags.api_override.update(abase.get_api_override_from_override_string(api_override))
          if api_override is not None and "Options" in api_override:
            symbols = convert_options_symbols(symbols, fsflags)
          print("symbols:" + str(symbols))
          param = [
              launch,
              copy.deepcopy(baseflags),
              copy.deepcopy(fsflags),
              list(symbols),
              out_dir,
              len(trading_dates) > 1,
              fs_dict['use_exchange_timestamp']
          ]
          params.append(param)
          cgroups.append(param)
      compare_groups[(trading_date, trange)] = cgroups

  if compare:
    nameparams = [
        flags.FLAGS.subsymbols,
        trading_date_given,
        flags.FLAGS.feed_machine,
        flags.FLAGS.compare_feed_machines,
        flags.FLAGS.compare_worker_ids,
        flags.FLAGS.compare_subsymbols,
        flags.FLAGS.compare_api_overrides,
        flags.FLAGS.postfix]
    filename = f"{combine_names(nameparams)}"

    def diff_ratio(numa, numb):
      return abs(float(numa) / (numa + numb) - 0.5) * 2

    compare_group_iter = iter(compare_groups.items())

    preruns = {}
    for compare_key, compare_params in compare_group_iter:
      for ith, param in enumerate(compare_params):
        preruns[(compare_key, ith)] = param[0](
            *param[1:], prerun=True, compare_idx=ith, rasterized=True)

    plt.rcParams['figure.figsize'] = 12, 11
    plt.rcParams['font.size'] = 9
    plt.rcParams['legend.fontsize'] = 9
    plt.rcParams['xtick.labelsize'] = 9
    plt.rcParams['ytick.labelsize'] = 9
    with PdfPages(f"{flags.FLAGS.out_dir}/{filename}.pdf") as pdf:
      for compare_key, compare_params in compare_groups.items():
        print(compare_key, compare_params)
        lines = []
        legends = []
        prev_nbookntrade = None
        savefig = False
        if flags.FLAGS.compare_use_exchange_timestamp is not None:
          savefig = True
        for ith, param in enumerate(compare_params):
          prerun = preruns[(compare_key, ith)]
          line_legends = prerun[0](**prerun[1])
          for line_legend_dbg in line_legends:
            lines.append(line_legend_dbg[0][0])
            legends.append(line_legend_dbg[1])
            if prev_nbookntrade is not None:
              ratio_diff = max(diff_ratio(prev_nbookntrade[0], line_legend_dbg[2][0]),
                               diff_ratio(prev_nbookntrade[1], line_legend_dbg[2][1]))
              if ratio_diff > 0.0001:
                savefig = True
            prev_nbookntrade = line_legend_dbg[2]
        if savefig:
          plt.legend(lines, legends, bbox_to_anchor=(1.0, -0.05), loc="upper right")
          plt.title(compare_key)
          plt.tight_layout()
          pdf.savefig()
        plt.close()
    return

  if len(params) <= 1:
    for param in params:
      param[0](*param[1:])
  else:
    outs = []
    for param in params:
      out = param[0](*param[1:], prerun=True, rasterized=True)
      outs.append(out)
    filename = f"{combine_names([flags.FLAGS.subsymbols, flags.FLAGS.api_override, flags.FLAGS.feed_machine, trading_date_given, flags.FLAGS.postfix])}"  # noqa: E501
    plt.rcParams['figure.figsize'] = 12, 8
    plt.rcParams['font.size'] = 9
    plt.rcParams['legend.fontsize'] = 9
    plt.rcParams['xtick.labelsize'] = 9
    plt.rcParams['ytick.labelsize'] = 9
    with PdfPages(f"{flags.FLAGS.out_dir}/{filename}.pdf") as pdf:
      for out in outs:
        out[0](**out[1])
        try:
          pdf.savefig()
        except Exception:
          pass
        plt.close()


# ./pyrunner coin/strategy/mm/tool/feed_archive_checker.py
#   --trading_date=20200101-20200102
#   --subsymbols=Okex:Futures:BTC-USD.QUARTER --worker_ids=0
#   --compare_feed_machines=feed-02.ap-northeast-1.aws,feed-02.cn-hongkong.aliyun
#   --split_hours=1
#   --use_feed_cache=True
if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('full_symbols', None, 'full symbol name')
  flags.DEFINE_string('compare_subsymbols', None, 'subscript symbol name')
  flags.DEFINE_string('subsymbols', None, 'subscript symbol name')
  flags.DEFINE_string('symbols', None, 'symbol name')
  flags.DEFINE_string('out_dir', 'pic', 'Output directory.')
  flags.DEFINE_string('postfix', None, 'filename postfix')
  # compare worker ids = ex) 1-2-1,2
  flags.DEFINE_string('compare_worker_ids', None, 'compare_worker_ids')
  # compare feed machines = ex) feed-01.ap-northeast-1.aws,feed-01.eu-west-1.aws, ...
  flags.DEFINE_string('compare_feed_machines', None, 'compare feed machines')
  flags.DEFINE_string('compare_api_overrides', None, 'compare api overrides')
  # --compare_feed_machine_api_overrides=feed-03.ap-northeast-1.aws^Futures:Huobi=,feed-05.ap-northeast-1.aws^Futures:Huobi=v1
  flags.DEFINE_string('compare_feed_machine_api_overrides', None, '')
  flags.DEFINE_string('compare_use_exchange_timestamp', None, '')
  flags.DEFINE_string('excel_file', None, '')
  flags.DEFINE_float('split_hours', None, '')
  flags.DEFINE_bool('print_all_feed', False, '')
  flags.DEFINE_integer('num_cpu', 12, '')
  flags.DEFINE_bool('use_exchange_timestamp', False, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  app.run(main)
